#!/usr/bin/env python3

import argparse
import f90nml
import netCDF4 as nc
import numpy as np
import os
import subprocess

parser = argparse.ArgumentParser('Check numerical accuracy order.')
parser.add_argument('-n', '--np', help='Processes to use for running tests', type=int, default=1)
args = parser.parse_args()

namelist = f90nml.namelist.Namelist()
namelist['hope_control'] = {
	'run_hours'        : 1,
	'history_interval' : 3600,
	'time_scheme'      : 'rk4',
	'recon_h_scheme'   : 'WENO2D',
	'recon_h_order'    : 3,
	'use_tci'          : False
}

exe = os.path.dirname(os.path.realpath(__file__)) + '/../build/hope.exe'

def run(cmd):
	print(f'==> {cmd}')
	res = subprocess.run(cmd, shell=True, check=True)

def mpiexec(namelist, args):
	run(f'mpiexec -np {args.np} {exe} {namelist}')

def run_hope(namelist):
	file_path = f'namelist.input'
	namelist.write(file_path, force=True)
	print(namelist)
	mpiexec(file_path, args)
	os.remove(file_path)

# ------------------------------------------------------------------------------
#                                     CASE 2
namelist['hope_control']['case_id'         ] = 2
namelist['hope_control']['case_name'       ] = 'hope_case2'
namelist['hope_control']['run_hours'       ] = 12
namelist['hope_control']['history_interval'] = 3600 * 12
namelist['hope_control']['recon_h_order'   ] = 3
configs = [
	{ 'nc': 45, 'dt': 400 },
	{ 'nc': 90, 'dt': 200 },
]

for config in configs:
	namelist['hope_control']['nc'] = config['nc']
	namelist['hope_control']['dt'] = config['dt']
	run_hope(namelist)

L1 = []; L2 = []; Linf = []; dx = []
for file_path in ('hope_case2.n45.nc', 'hope_case2.n90.nc'):
	f = nc.Dataset(file_path, 'r')
	J = f.variables['Jh'][:,0,:,:]
	h0 = f.variables['gh'][0,:]
	h1 = f.variables['gh'][-1,:]
	diff = np.abs(h1 - h0)
	L1.append(np.sum(J * diff) / np.sum(J * h0))
	L2.append(np.sum(J * diff**2) / np.sum(J * h0**2))
	Linf.append(np.max(diff) / np.max(h0))
	dx.append(f.dx)

order = {}
order['L1'] = -(np.log(L1[0]) - np.log(L1[1])) / (np.log(dx[0]) - np.log(dx[1]))
order['L2'] = -(np.log(L2[0]) - np.log(L2[1])) / (np.log(dx[0]) - np.log(dx[1]))
order['Linf'] = -(np.log(Linf[0]) - np.log(Linf[1])) / (np.log(dx[0]) - np.log(dx[1]))

print(L1)
print(L2)
print(Linf)
print(order)

# ------------------------------------------------------------------------------
#                                     CASE 5
namelist['hope_control']['case_id'         ] = 5
namelist['hope_control']['case_name'       ] = 'hope_case5'
namelist['hope_control']['run_hours'       ] = 20 * 24
namelist['hope_control']['history_interval'] = 86400
namelist['hope_control']['nc'              ] = 45
namelist['hope_control']['dt'              ] = 450
run_hope(namelist)

# ------------------------------------------------------------------------------
#                                     CASE 6
namelist['hope_control']['case_id'         ] = 6
namelist['hope_control']['case_name'       ] = 'hope_case6'
namelist['hope_control']['run_hours'       ] = 30 * 24
namelist['hope_control']['history_interval'] = 86400
namelist['hope_control']['nc'              ] = 45
namelist['hope_control']['dt'              ] = 450
run_hope(namelist)
